import os
import json
import threading
import queue
import tkinter as tk
from tkinter import scrolledtext, messagebox, simpledialog
import requests


class DeepSeekClientGUI:
    def __init__(self, root):
        self.root = root
        self.root.title("DeepSeek Chat Client")

        # 初始化API客户端
        self.client = None
        self.response_queue = queue.Queue()

        # 创建界面组件
        self.create_widgets()
        self.check_api_key()

    def create_widgets(self):
        """创建GUI组件"""
        # 对话历史展示区
        self.history_area = scrolledtext.ScrolledText(
            self.root,
            wrap=tk.WORD,
            width=60,
            height=20,
            font=("微软雅黑", 10)
        )
        self.history_area.pack(padx=10, pady=10)
        self.history_area.configure(state='disabled')

        # 用户输入区
        input_frame = tk.Frame(self.root)
        input_frame.pack(padx=10, pady=10, fill=tk.X)

        self.input_entry = tk.Entry(
            input_frame,
            font=("微软雅黑", 12),
            width=50
        )
        self.input_entry.pack(side=tk.LEFT, fill=tk.X, expand=True)
        self.input_entry.bind("<Return>", lambda event: self.send_message())

        self.send_button = tk.Button(
            input_frame,
            text="发送",
            command=self.send_message,
            font=("微软雅黑", 10)
        )
        self.send_button.pack(side=tk.LEFT, padx=5)
        self.status_label = tk.Label(text="⏳ 请求中...", fg="gray")

    def check_api_key(self):
        """检查API密钥"""
        api_key = os.getenv("DEEPSEEK_API_KEY")
        if not api_key:
            api_key = simpledialog.askstring(
                "API密钥",
                "请输入DeepSeek API密钥:",
                parent=self.root,
                show='*'
            )
        if api_key:
            try:
                self.client = DeepSeekClient(api_key)
            except Exception as e:
                messagebox.showerror("错误", str(e))
        else:
            messagebox.showwarning("警告", "未提供API密钥，程序将退出")
            self.root.destroy()

    def send_message(self):
        """发送用户消息"""
        user_input = self.input_entry.get().strip()
        if not user_input:
            return

        # 禁用输入框和按钮
        self.input_entry.config(state='disabled')
        self.send_button.config(state='disabled')

        # 显示用户消息
        self.update_history(f"用户: {user_input}\n\n", role="user")

        # 启动后台线程处理请求
        threading.Thread(
            target=self.handle_api_request,
            args=(user_input,),
            daemon=True
        ).start()

        # 启动队列监听
        self.root.after(100, self.process_response_queue)

    def handle_api_request(self, prompt):
        """处理API请求（优化版）"""
        try:
            # 移除冗余变量，直接调用方法
            self.client.chat(prompt, gui_queue=self.response_queue)
            self.response_queue.put(("[END]", ""))
        except Exception as e:
            self.response_queue.put(("[ERROR]", str(e)))

    def process_response_queue(self):
        """处理响应队列更新界面（完整修复版）"""
        while not self.response_queue.empty():
            msg_type, content = self.response_queue.get()

            if msg_type == "[ERROR]":
                messagebox.showerror("API错误", content)
            elif msg_type == "[DATA]":  # 新增分支
                #content = f"assistant: {content}\n\n"
                self.update_history(content, role="assistant")
            elif msg_type == "[END]":
                self.input_entry.config(state='normal')
                self.send_button.config(state='normal')
                self.input_entry.delete(0, tk.END)
                self.update_history("\n\n", role="system")

        self.root.after(100, self.process_response_queue)

    def update_history(self, content, role="assistant"):
        """更新对话历史"""
        self.history_area.configure(state='normal')

        # 添加不同角色的颜色标记
        if role == "user":
            self.history_area.tag_config("user", foreground="blue")
            self.history_area.insert(tk.END, content, "user")
        elif role == "assistant":
            self.history_area.tag_config("assistant", foreground="green")
            self.history_area.insert(tk.END, content, "assistant")
        else:
            self.history_area.insert(tk.END, content)

        self.history_area.see(tk.END)  # 自动滚动到底部
        self.history_area.configure(state='disabled')


class DeepSeekClient:
    """API客户端（与之前代码相同）"""

    def __init__(self, api_key):
        self.api_key = api_key
        self.base_url = "https://api.deepseek.com/v1/chat/completions"
        self.headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }

    def chat(self, prompt, gui_queue=None, model="deepseek-chat", max_tokens=2048):
        payload = {
            "model": model,
            "messages": [{"role": "user", "content": prompt}],
            "stream": True,
            "temperature": 0.7,
            "max_tokens": max_tokens
        }

        try:
            response = requests.post(
                self.base_url,
                headers=self.headers,
                json=payload,
                stream=True,
                timeout=30
            )
            response.raise_for_status()

            full_response = []
            for line in response.iter_lines():
                if line:
                    decoded_line = line.decode('utf-8')
                    if decoded_line.startswith("data: "):
                        chunk = decoded_line[6:]
                        if chunk == "[DONE]":
                            break
                        try:
                            data = json.loads(chunk)
                            content = data["choices"][0]["delta"].get("content", "")
                            full_response.append(content)
                            print(full_response)
                            if gui_queue:  # 关键修改点
                                gui_queue.put(("[DATA]", content))  # 明确推送数据
                        except json.JSONDecodeError:
                            continue

            return "".join(full_response)

        except requests.exceptions.RequestException as e:
            if gui_queue:
                gui_queue.put(("[ERROR]", str(e)))
            return None


if __name__ == "__main__":
    root = tk.Tk()
    app = DeepSeekClientGUI(root)
    root.mainloop()